package com.example.filter;

import lombok.Data;
import lombok.experimental.Accessors;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.time.Instant;
import java.util.HashMap;
import java.util.Map;

/**
 * @author lbq
 * @version 1.0
 * @date 2021/3/15 20:23
 */
@Component
public class IPFilterFactory extends AbstractGatewayFilterFactory<IPFilterFactory.Config> {

    private static final Map<String, Record> IP_MAP = new HashMap<>();

    private static final Long TIME = 60000L;
    private static final Integer NUM = 3;

    public IPFilterFactory() {
        super(Config.class);
    }

    @Override
    public GatewayFilter apply(Config config) {
        return this::filter;
    }

    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        String path = request.getURI().getPath();
        if (path.contains("/api/user/register")) {
            String host = request.getRemoteAddress().getHostName();
            Record record = IP_MAP.get(host);
            long now = Instant.now().toEpochMilli();
            if (record == null) {
                IP_MAP.put(host, new Record().setCount(1).setVisitTime(now));
            } else {
                if (now - record.getVisitTime() <= TIME && record.getCount() < NUM) {
                    IP_MAP.put(host, record.setCount(record.getCount() + 1));
                } else if (now - record.getVisitTime() <= TIME && record.getCount() >= NUM) {
                    ServerHttpResponse response = exchange.getResponse();
                    response.setStatusCode(HttpStatus.FORBIDDEN);
                    String data = "Request has been forbidden!";
                    DataBuffer wrap = response.bufferFactory().wrap(data.getBytes());
                    return response.writeWith(Mono.just(wrap));
                } else if (now - record.getVisitTime() > TIME) {
                    IP_MAP.put(host, record.setCount(1));
                }
            }
        }
        return chain.filter(exchange);
    }

    @Override
    public String name() {
        return "IPFilter";
    }

    public static class Config {
    }

    @Data
    @Accessors(chain = true)
    public static class Record {
        // 访问次数
        private Integer count;
        // 第一次访问时间或上一分钟第一次访问时间
        private Long visitTime;
    }
}
